/**
 * Copyright (c) 2025 QINGMAO INTELLIGENCE TECHNOLOGY (BEIJING) CO., LTD. and Huawei Technologies Co., Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include "config.h"
#include "hccl/hccl.h"
#include <mpi.h>

enum class CommBackend
{
    MPI,
    HCCL
};

class CommDomainManager
{
public:
    static void Initialize(ParallelConfig& config, MPI_Comm workerGlobalComm = MPI_COMM_WORLD);

    static void Finalize();

    enum class CommType
    {
        GLOBAL,
        TP,
        EP,
        DP,
    };

    static HcclComm& GetHcclComm(CommType type)
    {
        switch (type)
        {
            case CommType::GLOBAL: return mGlobalCommHccl;
            case CommType::TP: return mTpCommHccl;
            case CommType::EP: return mEpCommHccl;
            case CommType::DP: return mDpCommHccl;
        }
        return mGlobalCommHccl;
    }

    static MPI_Comm& GetMPIComm(CommType type)
    {
        switch (type)
        {
        case CommType::GLOBAL: return mGlobalCommMpi;
        case CommType::TP: return mTpCommMpi;
        case CommType::EP: return mEpCommMpi;
        case CommType::DP: return mDpCommMpi;
        }
        return mGlobalCommMpi;
    }

private:
    static MPI_Comm mGlobalCommMpi, mTpCommMpi, mEpCommMpi, mDpCommMpi;
    static HcclComm mGlobalCommHccl, mTpCommHccl, mEpCommHccl, mDpCommHccl;
};
